import datetime
import random
import math
import numpy

import common
import utils
import warningExtractor
import stormReportExtractor
import radarExtractor
import radarAnalysis


class VortTrack:
	def __init__(self, initFrame):
		self.frames = [initFrame]
		
		self.avgHdg = 0
		self.avgSpd = 0
		self.avgGrad = 0
		self.avgRef = 0
		self.avgVar = 0
		self.avgVelDiff = 0
		
		self.avgScore = 0
		self.confidence = 0
		self.center = (0, 0)
		
	def addFrame(self, frame):
		self.frames.append(frame)
		
		# compute track linear regression
		xi = [f.shearFeature.centerLatLon[1] for f in self.frames]
		A = numpy.array([xi, numpy.ones(len(self.frames))])
		y = [f.shearFeature.centerLatLon[0] for f in self.frames]
		w = numpy.linalg.lstsq(A.T,y)[0]
		
		xChg = self.frames[-1].shearFeature.centerLatLon[1]-self.frames[0].shearFeature.centerLatLon[1]
		yChg = self.frames[-1].shearFeature.centerLatLon[0]-self.frames[0].shearFeature.centerLatLon[0]
		self.avgHdg = utils.hdgSlope(w[0], xChg, yChg)
		self.avgSpd = float(sum([f.frameSpd for f in self.frames]))/(len(self.frames)-1)
		
		# compute averages
		self.avgRef = float(sum([x.shearFeature.avgRef for x in self.frames]))/(len(self.frames))
		self.avgGrad = float(sum([x.shearFeature.avgVelGrad for x in self.frames]))/(len(self.frames))
		self.avgVar = float(sum([x.shearFeature.avgVar for x in self.frames]))/(len(self.frames))
		self.avgVelDiff = float(sum([x.shearFeature.avgVelDiff for x in self.frames]))/(len(self.frames))
		self.avgScore = float(sum([x.frameScore for x in self.frames]))/len(self.frames)
		cX = float(sum([x.shearFeature.center[0] for x in self.frames]))/len(self.frames)
		cY = float(sum([x.shearFeature.center[1] for x in self.frames]))/len(self.frames)
		self.center = (cX, cY)
			

class VortTrackFrame:
	
	MaxDistPerFrame = 6.0
	
	def __init__(self, feature):
		self.shearFeature = feature
		self.frameHdg = 0
		self.frameSpd = 0
		self.frameDist = 0
		self.frameScore = 0
		
		self.imgId = -1
		self.segmentId = -1
	
	def update(self, hdg, lastHdg, spd, lastSpd, dist):
		""" score the current frame based on the feature, hdg, spd, and dist. """
		self.frameHdg = hdg
		self.frameSpd = spd
		self.frameDist = dist
		if lastHdg > 0 and lastSpd > 0:
			self.frameScore = (1.0 - abs(utils.hdgDiff(hdg, lastHdg))/180.0) + \
							  (1.0 - abs(spd - lastSpd)/50.0) + \
							  (1.0 - dist/VortTrackFrame.MaxDistPerFrame) + \
							  (self.shearFeature.score / 100.0)
		else:
			self.frameScore = (1.0 - float(dist)/VortTrackFrame.MaxDistPerFrame) + \
							  (self.shearFeature.score / 100.0)

def vortTracker(warn, verbose = False):
	""" Finds radar image at the issue time of warn. Uses the image, the one before it, and 
		the one after it to produce a list of (heading, speed) pairs indicating 
		the tracks of all of the velocity couplets in the image."""
	
	if verbose:
		print 'warning: %s county, %s'%(warn["countyName"], str(warn["iTime"]))
	
	# find x/y coords of warning lat/lon
	countyLatLon = warn["countyLatLon"]
	countyXYCoords = utils.latLonToXY(countyLatLon[0], countyLatLon[1])
	
	# frames before the warning time
	bTrack = 7
	# frames ahead of the warning time
	fTrack = 0
	
	stormRel0 = radarExtractor.getClosestRadarData(warn["iTime"], bNum = bTrack, fNum = fTrack, product='stormRel0')
	stormRel1 = radarExtractor.getClosestRadarData(warn["iTime"], bNum = bTrack, fNum = fTrack, product='stormRel1')
	stormRel2 = radarExtractor.getClosestRadarData(warn["iTime"], bNum = bTrack, fNum = fTrack, product='stormRel2')
	stormRel3 = radarExtractor.getClosestRadarData(warn["iTime"], bNum = bTrack, fNum = fTrack, product='stormRel3')
	baseRef = radarExtractor.getClosestRadarData(warn["iTime"], bNum = bTrack, fNum = fTrack, product='baseRef0')

	# convert the chosen radar files to NetCDF
	stormRel0Nc = radarExtractor.radarToNetCDF(stormRel0)
	stormRel1Nc = radarExtractor.radarToNetCDF(stormRel1)
	stormRel2Nc = radarExtractor.radarToNetCDF(stormRel2)
	stormRel3Nc = radarExtractor.radarToNetCDF(stormRel3)
	baseRefNc = radarExtractor.radarToNetCDF(baseRef)

	if len(stormRel0Nc) == 0 or len(stormRel1Nc) == 0 or \
	   len(stormRel2Nc) == 0 or len(stormRel3Nc) == 0 or len(baseRefNc) == 0: 
		raise Exception("vortTracker()", "Failed to gather NC file data")

	# extract the radar data from the NetCDF files
	baseRefData = []
	stormRel0Data = []
	stormRel1Data = []
	stormRel2Data = []
	stormRel3Data = []
	
	if not (len(baseRefNc) == len(stormRel0Nc) == len(stormRel1Nc) == \
			len(stormRel2Nc) == len(stormRel3Nc)):
				raise Exception("vortTester", "NC file lists are not the same length")
	
	for i in range(len(baseRefNc)):
		data, prefix = radarExtractor.getRadarData(baseRefNc[i], flat = False, flip = True)
		baseRefData.append(data)
		
		data, prefix = radarExtractor.getRadarData(stormRel0Nc[i], flat = False, flip = True)
		stormRel0Data.append(data)
		
		data, prefix = radarExtractor.getRadarData(stormRel1Nc[i], flat = False, flip = True)
		stormRel1Data.append(data)
		
		data, prefix = radarExtractor.getRadarData(stormRel2Nc[i], flat = False, flip = True)
		stormRel2Data.append(data)
		
		data, prefix = radarExtractor.getRadarData(stormRel3Nc[i], flat = False, flip = True)
		stormRel3Data.append(data)

	shearFeatures = []
	
	for i in range(len(baseRefData)):
		if verbose:
			print 'finding shear features for image %d'%i
		curShearFeatures = radarAnalysis.findShearFeatures({'stormRel0':stormRel0Data[i], 'stormRel1':stormRel1Data[i], 
										   'stormRel2':stormRel2Data[i], 'stormRel3':stormRel3Data[i]}, 
										   baseRefData[i], targetLatLon = countyLatLon, maxDist = 30, verbose=True)
		
		shearFeatures.append(curShearFeatures)
		
		if verbose:
			for k in range(len(curShearFeatures)):
				c = utils.xyToLatLon(curShearFeatures[k].center[0], curShearFeatures[k].center[1])
				print '(%.2f, %.2f): grad = %.2f, ref = %.2f, dist = %.2f, var = %.2f, score = %f.2'%(c[0], c[1],
																			curShearFeatures[k].avgVelGrad, 
																			curShearFeatures[k].avgRef, 
																			curShearFeatures[k].targetDist,
																			curShearFeatures[k].avgVar,
																			curShearFeatures[k].score)
			
		
			print 'found %d shear features...'%len(curShearFeatures)
	
	# find recurring features
	vortTracks = findVortTracks(shearFeatures, verbose)
	if verbose:
		print '--vort tracks--'
		for key in vortTracks.keys():
			print '%d : %d'%(key, len(vortTracks[key].frames))
			for frame in vortTracks[key].frames:
				centerLatLon = utils.xyToLatLon(frame.shearFeature.center[0], frame.shearFeature.center[1])
				print '%.3f, %.3f: grad = %.3f, ref = %.3f, var = %.3f, dist = %.3f, score = %.3f'%(centerLatLon[0], centerLatLon[1], 
																					  frame.shearFeature.avgVelGrad, frame.shearFeature.avgRef,
																					  frame.shearFeature.avgVar, frame.shearFeature.targetDist, 
																					  frame.shearFeature.score)
			c = utils.xyToLatLon(vortTracks[key].center[0], vortTracks[key].center[1])
			print 'hdg = %.3f, spd = %.3f, score = %.3f, confidence = %.3f, center = (%.3f, %.3f)'%(vortTracks[key].avgHdg, vortTracks[key].avgSpd, 
																									vortTracks[key].avgScore, vortTracks[key].confidence, 
																									c[0], c[1])		
			print
	return vortTracks

def findVortTracks(shearFeatures, verbose = False):
	# max dist in mi that a couplet can move per radar frame to be considered a track. 
	maxDistPerFrame = 12.0
	
	# max change in storm heading per frame in degrees
	maxTrackHdgChange = 120.0
	
	# min/max average track speeds (mph)
	minTrackSpd = 2.0
	maxTrackSpd = 100.0
	
	# min distance (mi) that tracks must be apart
	minTrackSeparation = 7.0
	
	# min number of shear features in a track
	minTrackLength = 4
	
	# max score difference between two tracks
	maxScoreDiff = 60.0
	
	vortTracks = {}
	maxVortTrackId = 0
	
	def findNextTrackFrame(i, k, frameJump = 0):
		""" find the best choice for the next couplet in the current track """
		
		potentialTrackSteps = []
		
		def nonNegScore():
			for step in potentialTrackSteps:
				if step.frameScore > 0:
					return True
			return False
		
		h = i + frameJump + 1
		initH = h
		
		while not nonNegScore() and h < len(shearFeatures):
			for index in range(len(shearFeatures[h])):
				curFeature = shearFeatures[h][index]
				frame = VortTrackFrame(curFeature)
				
				# find the lat/lon coords of the current feature and the closest feature in the next frame
				pt1LatLon = utils.xyToLatLon(shearFeatures[i][k].center[0], shearFeatures[i][k].center[1])
				pt2LatLon = utils.xyToLatLon(curFeature.center[0], curFeature.center[1])

				# calculate the distance
				frameDist = utils.latLonDist(pt1LatLon, pt2LatLon)
				# calculate the heading
				frameHdg = utils.latLonFlatHeading(pt1LatLon, pt2LatLon)
				# calculate the speed -- ASSUME EACH FRAME IS 4 MINUTES
				frameSpd = frameDist / ((h-i)*4/60.0)
				
				trackHdg = -1
				trackSpd = -1
				curTrackId = shearFeatures[i][k].trackId
				if len(vortTracks[curTrackId].frames) > 1:
					trackHdg = vortTracks[curTrackId].avgHdg
					trackSpd = vortTracks[curTrackId].avgSpd
				
				frame.update(frameHdg, trackHdg, frameSpd, trackSpd, frameDist)
				frame.frameScore *= (0.5)**(h-initH)
				frame.imgId = h
				frame.segmentId = index
				
				if verbose:
					print '(%d, %d) -> (%d, %d): (%.2f, %.2f) -> (%.2f, %.2f): %.2f mi, hdgDiff = %.2f, spdDiff = %.2f, score = %.4f'%(i, k, h, index,
																				pt1LatLon[0], pt1LatLon[1],
																				pt2LatLon[0], pt2LatLon[1], frameDist, 
																				abs(utils.hdgDiff(frameHdg, trackHdg)), abs(frameSpd - trackSpd),
																				frame.frameScore)
				
				potentialTrackSteps.append(frame)
			h += 1
		
		potentialTrackSteps.sort(key = lambda x:x.frameScore, reverse = True)
		return potentialTrackSteps

	# loop through the list of radar images
	for i in range(len(shearFeatures)):
		
		# loop through the shear features in the current image
		for k in range(len(shearFeatures[i])):
			
			# if the current feature doesn't have an id, make one. each feature will be labeled at its first appearance
			if shearFeatures[i][k].trackId == -1:
				shearFeatures[i][k].trackId = maxVortTrackId
				vortTracks[maxVortTrackId] = VortTrack(VortTrackFrame(shearFeatures[i][k]))
				maxVortTrackId += 1
			
			curTrackId = shearFeatures[i][k].trackId
			if verbose:
				print 'curTrack = %d'%curTrackId
			frameChoices = findNextTrackFrame(i, k)
			
			# add the highest scoring next frame choice
			if len(frameChoices) > 0:
				topChoiceFrame = frameChoices[0]
				
				if topChoiceFrame.frameDist < maxDistPerFrame:
					vortTracks[curTrackId].addFrame(topChoiceFrame)
					shearFeatures[topChoiceFrame.imgId][topChoiceFrame.segmentId].trackId = curTrackId	
					
					if verbose:
						print '%d: (%d, %d)->(%d, %d): dist = %.2f, frame heading = %.2f, ' \
							  'track heading = %.2f, frame speed = %.2f, track speed = %.2f'%(curTrackId, i, k, 
																							  topChoiceFrame.imgId, topChoiceFrame.segmentId, 
																							  frameChoices[0].frameDist, frameChoices[0].frameHdg, 
																							  vortTracks[curTrackId].avgHdg, 
																							  frameChoices[0].frameSpd, vortTracks[curTrackId].avgSpd)
						print
				
			
	finalVortTracks = {}
	
	# purge bad tracks
	for key in vortTracks.keys():
		
		# search for duplicate tracks and tracks that have wildly different scores
		duplicateTrack = False
		largeScoreDiff = False
		for key2 in vortTracks.keys():
			if key == key2: continue
			c1 = utils.xyToLatLon(vortTracks[key].center[0], vortTracks[key].center[1])
			c2 = utils.xyToLatLon(vortTracks[key2].center[0], vortTracks[key2].center[1])
			cDist = utils.latLonDist(c1, c2)
			# if tracks too close, don't include the lower-scoring one
			if cDist < minTrackSeparation:
				if vortTracks[key].avgScore < vortTracks[key2].avgScore:
					duplicateTrack = True
					break
			# if scores are very different, don't include the lower scoring one
			if vortTracks[key2].avgScore - vortTracks[key].avgScore > maxScoreDiff:
				#largeScoreDiff = True
				break
		
		# don't include tracks that are too short, too fast, too slow, are duplicates, or have a wildly lower score than another track
		if len(vortTracks[key].frames) < minTrackLength or \
		   vortTracks[key].avgSpd > maxTrackSpd or vortTracks[key].avgSpd < minTrackSpd or \
		   duplicateTrack or largeScoreDiff:
			continue
		else:
			finalVortTracks[key] = vortTracks[key]

	return finalVortTracks

def findClosestCouplets(targetShearSegment, shearSegmentList):
	"""Takes in a shear feature, target, and a list of shear features,
		and returns a list of indicies in the list of shear features, ranking them
		from closest to farthest."""
	targetLatLon = utils.xyToLatLon(targetShearSegment.center[0], targetShearSegment.center[1])
	ranking = []
	minDist = -1
	minIndex = -1
	for i in range(len(shearSegmentList)):
		curCenter = shearSegmentList[i].center
		curDist = utils.latLonDist(utils.xyToLatLon(targetShearSegment.center[0], targetShearSegment.center[1]), utils.xyToLatLon(curCenter[0], curCenter[1]))
		ranking.append((i, curDist))
	
	ranking.sort(key = lambda x:x[1], reverse = False)
	return ranking

def vortTrackerTest():
	print "-------- vortTracker() test ----------"
	
	print 'loading warnings...'
	warnings = warningExtractor.loadTorWarnings(range(2005,2010+1))

	maxDist = 80
	minDist = 5
	distValidWarnings = []

	for warn in warnings:
		dist = utils.latLonDist(warn["countyLatLon"], common.KTLXLocation)
		if dist <= maxDist and dist >= minDist:
			distValidWarnings.append(warn)

	print '%d dist valid warnings'%len(distValidWarnings)
	warn = distValidWarnings[45]
	vortTracker(warn, verbose = True)
	
	print 'tracked couplets for: \nwarning for %s, %s'%(warn["countyName"], str(warn["iTime"]))

